"""
Domain decomposition module for locELM.

This module handles domain partitioning and collocation point generation
for the local Extreme Learning Machine method.
"""

import numpy as np
from typing import Tuple, List, Optional
from scipy.special import roots_legendre


class Domain:
    """Base class for domain representation."""

    def __init__(self, bounds: List[Tuple[float, float]], n_subdomains: List[int]):
        """
        Initialize domain.

        Parameters
        ----------
        bounds : List[Tuple[float, float]]
            Domain bounds for each dimension [(a1, b1), (a2, b2), ...]
        n_subdomains : List[int]
            Number of subdomains in each dimension
        """
        self.bounds = bounds
        self.n_subdomains = n_subdomains
        self.ndim = len(bounds)
        self.n_total = np.prod(n_subdomains)

        # Generate subdomain boundaries
        self.subdomain_bounds = self._generate_subdomain_bounds()

    def _generate_subdomain_bounds(self) -> List[np.ndarray]:
        """Generate subdomain boundary coordinates for each dimension."""
        bounds_list = []
        for (a, b), n in zip(self.bounds, self.n_subdomains):
            bounds_list.append(np.linspace(a, b, n + 1))
        return bounds_list

    def get_subdomain_region(self, subdomain_idx: int) -> List[Tuple[float, float]]:
        """
        Get the bounds of a specific subdomain.

        Parameters
        ----------
        subdomain_idx : int
            Linear index of subdomain

        Returns
        -------
        List[Tuple[float, float]]
            Bounds for each dimension
        """
        multi_idx = self._linear_to_multi_index(subdomain_idx)
        region = []
        for dim, idx in enumerate(multi_idx):
            a = self.subdomain_bounds[dim][idx]
            b = self.subdomain_bounds[dim][idx + 1]
            region.append((a, b))
        return region

    def _linear_to_multi_index(self, linear_idx: int) -> Tuple[int, ...]:
        """Convert linear subdomain index to multi-dimensional index."""
        if self.ndim == 1:
            return (linear_idx,)
        elif self.ndim == 2:
            ny = self.n_subdomains[1]
            m = linear_idx // ny
            n = linear_idx % ny
            return (m, n)
        elif self.ndim == 3:
            ny = self.n_subdomains[1]
            nz = self.n_subdomains[2]
            m = linear_idx // (ny * nz)
            n = (linear_idx % (ny * nz)) // nz
            l = linear_idx % nz
            return (m, n, l)
        else:
            raise NotImplementedError("Only up to 3D supported")

    def _multi_to_linear_index(self, multi_idx: Tuple[int, ...]) -> int:
        """Convert multi-dimensional subdomain index to linear index."""
        if self.ndim == 1:
            return multi_idx[0]
        elif self.ndim == 2:
            m, n = multi_idx
            ny = self.n_subdomains[1]
            return m * ny + n
        elif self.ndim == 3:
            m, n, l = multi_idx
            ny = self.n_subdomains[1]
            nz = self.n_subdomains[2]
            return m * ny * nz + n * nz + l
        else:
            raise NotImplementedError("Only up to 3D supported")

    def find_subdomain(self, point: np.ndarray) -> int:
        """
        Find which subdomain a point belongs to.

        Parameters
        ----------
        point : np.ndarray
            Point coordinates (1D array of length ndim)

        Returns
        -------
        int
            Linear index of subdomain containing the point
        """
        multi_idx = []
        for dim in range(self.ndim):
            coord = point[dim]
            bounds = self.subdomain_bounds[dim]

            # Find which interval the coordinate falls in
            for i in range(len(bounds) - 1):
                if bounds[i] <= coord <= bounds[i + 1]:
                    multi_idx.append(i)
                    break
            else:
                # Point is outside domain, use closest subdomain
                if coord < bounds[0]:
                    multi_idx.append(0)
                else:
                    multi_idx.append(len(bounds) - 2)

        return self._multi_to_linear_index(tuple(multi_idx))


class CollocationPoints:
    """Generate collocation points in subdomains."""

    @staticmethod
    def uniform(region: List[Tuple[float, float]],
                n_points: List[int]) -> np.ndarray:
        """
        Generate uniform grid points in a region.

        Parameters
        ----------
        region : List[Tuple[float, float]]
            Bounds for each dimension
        n_points : List[int]
            Number of points in each dimension

        Returns
        -------
        np.ndarray
            Collocation points of shape (n_total, ndim)
        """
        grids = []
        for (a, b), n in zip(region, n_points):
            grids.append(np.linspace(a, b, n))

        # Create meshgrid and flatten
        mesh = np.meshgrid(*grids, indexing='ij')
        points = np.stack([m.flatten() for m in mesh], axis=1)
        return points

    @staticmethod
    def gauss_lobatto_legendre(region: List[Tuple[float, float]],
                               n_points: List[int]) -> np.ndarray:
        """
        Generate Gauss-Lobatto-Legendre quadrature points.

        Parameters
        ----------
        region : List[Tuple[float, float]]
            Bounds for each dimension
        n_points : List[int]
            Number of points in each dimension

        Returns
        -------
        np.ndarray
            Collocation points of shape (n_total, ndim)
        """
        grids = []
        for (a, b), n in zip(region, n_points):
            # Get Gauss-Lobatto-Legendre points on [-1, 1]
            if n == 1:
                gll_points = np.array([0.0])
            elif n == 2:
                gll_points = np.array([-1.0, 1.0])
            else:
                # Use Legendre polynomial roots for interior points
                interior, _ = roots_legendre(n - 2)
                gll_points = np.concatenate([[-1.0], interior, [1.0]])

            # Map to [a, b]
            points = 0.5 * (b - a) * (gll_points + 1) + a
            grids.append(points)

        # Create meshgrid and flatten
        mesh = np.meshgrid(*grids, indexing='ij')
        points = np.stack([m.flatten() for m in mesh], axis=1)
        return points

    @staticmethod
    def random(region: List[Tuple[float, float]],
               n_points: int,
               n_boundary: Optional[int] = None,
               seed: Optional[int] = None) -> np.ndarray:
        """
        Generate random collocation points.

        Parameters
        ----------
        region : List[Tuple[float, float]]
            Bounds for each dimension
        n_points : int
            Total number of points
        n_boundary : Optional[int]
            Number of points on boundary (if None, computed automatically)
        seed : Optional[int]
            Random seed

        Returns
        -------
        np.ndarray
            Collocation points of shape (n_points, ndim)
        """
        if seed is not None:
            np.random.seed(seed)

        ndim = len(region)
        points = []

        # Generate interior points
        n_interior = n_points - (n_boundary or 0)
        interior_points = []
        for a, b in region:
            interior_points.append(np.random.uniform(a, b, n_interior))
        points.append(np.stack(interior_points, axis=1))

        # Generate boundary points if specified
        if n_boundary:
            # Distribute boundary points across all boundaries
            for dim in range(ndim):
                n_per_face = n_boundary // (2 * ndim)
                for boundary_val in [region[dim][0], region[dim][1]]:
                    boundary_points = []
                    for d, (a, b) in enumerate(region):
                        if d == dim:
                            boundary_points.append(np.full(n_per_face, boundary_val))
                        else:
                            boundary_points.append(np.random.uniform(a, b, n_per_face))
                    points.append(np.stack(boundary_points, axis=1))

        return np.vstack(points)


class SubdomainInterface:
    """Manages interfaces between subdomains for continuity conditions."""

    def __init__(self, domain: Domain):
        """Initialize subdomain interface manager."""
        self.domain = domain
        self.interfaces = self._find_interfaces()

    def _find_interfaces(self) -> List[dict]:
        """Find all subdomain interfaces."""
        interfaces = []

        for idx in range(self.domain.n_total):
            multi_idx = self.domain._linear_to_multi_index(idx)

            # Check each dimension for adjacent subdomains
            for dim in range(self.domain.ndim):
                # Check positive direction
                neighbor_idx = list(multi_idx)
                neighbor_idx[dim] += 1

                if neighbor_idx[dim] < self.domain.n_subdomains[dim]:
                    neighbor_linear = self.domain._multi_to_linear_index(tuple(neighbor_idx))

                    # Interface coordinate
                    interface_coord = self.domain.subdomain_bounds[dim][multi_idx[dim] + 1]

                    interfaces.append({
                        'subdomain1': idx,
                        'subdomain2': neighbor_linear,
                        'dimension': dim,
                        'coordinate': interface_coord,
                        'type': 'interior'
                    })

        return interfaces

    def get_boundary_interfaces(self, subdomain_idx: int) -> List[dict]:
        """Get boundary interfaces for a specific subdomain."""
        multi_idx = self.domain._linear_to_multi_index(subdomain_idx)
        boundaries = []

        for dim in range(self.domain.ndim):
            # Check if on lower boundary
            if multi_idx[dim] == 0:
                boundaries.append({
                    'subdomain': subdomain_idx,
                    'dimension': dim,
                    'coordinate': self.domain.subdomain_bounds[dim][0],
                    'type': 'lower'
                })

            # Check if on upper boundary
            if multi_idx[dim] == self.domain.n_subdomains[dim] - 1:
                boundaries.append({
                    'subdomain': subdomain_idx,
                    'dimension': dim,
                    'coordinate': self.domain.subdomain_bounds[dim][-1],
                    'type': 'upper'
                })

        return boundaries
